import os
from dotenv import load_dotenv
from openai import OpenAI
from langchain.chains import RetrievalQA
from langchain_community.vectorstores import FAISS
from langchain_community.embeddings import DashScopeEmbeddings
from langchain.docstore.document import Document
from typing import List, Dict, Any

# 加载环境变量
load_dotenv()

class LLMWrapper:
    def __init__(self):
        # 通义千问API配置
        self.api_key = os.getenv('TONGYI_API_KEY', 'sk-******')
        self.api_base_url = os.getenv('TONGYI_API_BASE_URL', 'https://dashscope.aliyuncs.com/compatible-mode/v1')
        
        # 初始化OpenAI客户端（用于调用通义千问API）
        try:
            self.client = OpenAI(
                api_key=self.api_key,
                base_url=self.api_base_url
            )
        except TypeError as e:
            # 处理旧版本openai库可能不支持某些参数的情况
            if "proxies" in str(e):
                self.client = OpenAI(
                    api_key=self.api_key,
                    base_url=self.api_base_url,
                    http_client=None
                )
            else:
                raise e
        
        # 初始化嵌入模型
        try:
            self.embeddings = DashScopeEmbeddings(
                model="text-embedding-v1",
                dashscope_api_key=self.api_key
            )
        except TypeError as e:
            # 处理旧版本库可能不支持某些参数的情况
            if "proxies" in str(e):
                self.embeddings = DashScopeEmbeddings(
                    model="text-embedding-v1",
                    dashscope_api_key=self.api_key
                )
            else:
                raise e
        
        # 初始化向量存储
        self.vector_store = None
        
    def generate_text(self, prompt: str, model: str = "qwen-plus", max_tokens: int = 1000, temperature: float = 0.7) -> str:
        """生成文本回复"""
        try:
            response = self.client.chat.completions.create(
                model=model,
                messages=[
                    {"role": "system", "content": "你是一个基于知识图谱的智能问答助手，可以回答关于资源关系网络的各种问题。"},
                    {"role": "user", "content": prompt}
                ],
                max_tokens=max_tokens,
                temperature=temperature
            )
            return response.choices[0].message.content
        except Exception as e:
            print(f"LLM调用失败: {str(e)}")
            return f"抱歉，我暂时无法回答这个问题。错误信息: {str(e)}"
    
    def build_vector_store(self, documents: List[Document]) -> None:
        """构建向量存储"""
        if not documents:
            # 如果没有文档，创建一个空的向量存储
            dummy_document = Document(page_content="placeholder", metadata={"source": "dummy"})
            self.vector_store = FAISS.from_documents([dummy_document], self.embeddings)
        else:
            self.vector_store = FAISS.from_documents(documents, self.embeddings)
    
    def search_similar_documents(self, query: str, k: int = 5) -> List[Document]:
        """搜索与查询相似的文档"""
        if not self.vector_store:
            return []
        
        results = self.vector_store.similarity_search(query, k=k)
        return results
    
    def rag_qa(self, query: str, context_docs: List[Document], model: str = "qwen-plus") -> Dict[str, Any]:
        """使用RAG技术回答问题"""
        # 构建向量存储
        self.build_vector_store(context_docs)
        
        # 准备上下文
        context_text = "\n".join([f"[{i+1}] {doc.page_content}" for i, doc in enumerate(context_docs)])
        
        # 构建提示
        prompt = f"""
你需要基于以下提供的知识（用[1]、[2]等标记的内容）来回答用户的问题。
如果提供的知识不足以回答问题，请直接说你无法回答，不要编造信息。
请确保你的回答准确、简明，并引用相关的知识来源（如[1]、[2]等）。

知识：
{context_text}

用户问题：{query}

回答：
"""
        
        # 调用大模型生成回答
        answer = self.generate_text(prompt, model=model)
        
        # 返回回答和相关文档
        return {
            "answer": answer,
            "context_docs": context_docs
        }
    
    def complex_query_analysis(self, query: str) -> Dict[str, Any]:
        """分析复杂查询，提取实体和关系"""
        # 使用大模型分析查询
        analysis_prompt = f"""
请分析以下查询，提取其中的实体、关系和条件，并以JSON格式返回：
{{
  "entities": ["实体1", "实体2"...],
  "relationships": ["关系1", "关系2"...],
  "conditions": ["条件1", "条件2"...],
  "query_type": "查询类型（如：单实体查询、多实体关系查询、条件过滤查询等）"
}}

查询：{query}
"""
        
        analysis_result = self.generate_text(analysis_prompt, model="qwen-plus")
        
        try:
            import json
            return json.loads(analysis_result)
        except Exception as e:
            print(f"查询分析结果解析失败: {str(e)}")
            return {
                "entities": [],
                "relationships": [],
                "conditions": [],
                "query_type": "未知",
                "raw_analysis": analysis_result
            }